# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt

class DrawTreeUtil:
    def __init__(self):
        self.decisionNode = dict(boxstyle="sawtooth", fc="0.8")
        self.leafNode = dict(boxstyle="round4", fc="0.8")
        self.arrow_args = dict(arrowstyle="<-")


    '''
    递归获取叶节点的数目
    @return 叶节点的数目
    '''
    def getNumLeafs(self, myTree):
        numLeafs = 0
        firstStr = list(myTree.keys())[0]
        secondDict = myTree[firstStr]
        for key in secondDict.keys():
            if type(secondDict[key]).__name__ == 'dict':
                # 如果节点是字典类型则递归
                numLeafs += self.getNumLeafs(secondDict[key])
            else:
                numLeafs += 1
        return numLeafs

    '''
    递归获取树的层数
    @return 树的层数
    '''
    def getTreeDepth(self, myTree):
        maxDepth = 0
        firstStr = list(myTree.keys())[0]
        secondDict = myTree[firstStr]
        for key in secondDict.keys():
            if type(secondDict[key]).__name__ == 'dict':
                thisDepth = 1 + self.getTreeDepth(secondDict[key])
            else:
                thisDepth = 1
            if thisDepth > maxDepth:
                maxDepth = thisDepth
        return maxDepth

    def plotNode(self, nodeTxt, centerPt, parentPt, nodeType):
        self.createPlot_ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                                xytext=centerPt, textcoords='axes fraction',
                                va="center", ha="center", bbox=nodeType, arrowprops=self.arrow_args)

    def plotMidText(self, cntrPt, parentPt, txtString):
        xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
        yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
        self.createPlot_ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)

    def plotTree(self, myTree, parentPt, nodeTxt):
        numLeafs = self.getNumLeafs(myTree)
        depth = self.getTreeDepth(myTree)
        firstStr = list(myTree.keys())[0]
        cntrPt = (self.plotTree_xOff + (1.0 + float(numLeafs)) / 2.0 / self.plotTree_totalW, self.plotTree_yOff)
        self.plotMidText(cntrPt, parentPt, nodeTxt)
        self.plotNode(firstStr, cntrPt, parentPt, self.decisionNode)
        secondDict = myTree[firstStr]
        self.plotTree_yOff = self.plotTree_yOff - 1.0 / self.plotTree_totalD
        for key in secondDict.keys():
            if type(secondDict[key]).__name__ == 'dict':
                self.plotTree(secondDict[key], cntrPt, str(key))
            else:
                self.plotTree_xOff = self.plotTree_xOff + 1.0 / self.plotTree_totalW
                self.plotNode(secondDict[key], (self.plotTree_xOff, self.plotTree_yOff), cntrPt, self.leafNode)
                self.plotMidText((self.plotTree_xOff, self.plotTree_yOff), cntrPt, str(key))
        self.plotTree_yOff = self.plotTree_yOff + 1.0 / self.plotTree_totalD

    def createPlot(self, inTree):
        fig = plt.figure(1, facecolor='white')
        fig.clf()
        axprops = dict(xticks=[], yticks=[])
        self.createPlot_ax1 = plt.subplot(111, frameon=False, **axprops)  # no ticks
        self.plotTree_totalW = float(self.getNumLeafs(inTree))
        self.plotTree_totalD = float(self.getTreeDepth(inTree))
        self.plotTree_xOff = -0.5 / self.plotTree_totalW
        self.plotTree_yOff = 1.0
        self.plotTree(inTree, (0.5, 1.0), '')
        plt.show()